{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "guide_for_new_tasks.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "JndnmDMp66FL",
        "hmLxdk2D5YQG"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL"
      },
      "source": [
        "##### Copyright 2020 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hMqWDc_m6rUC"
      },
      "source": [
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vsS6OoYqWfqb"
      },
      "source": [
        "\n",
        "# How to use the GFSA layer for new tasks\n",
        "\n",
        "This notebook describes the high-level process for using the GFSA layer in a new task, specifically focusing on how to represent a new type of graph as an MDP so that you can use the GFSA layer.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wqAn-pTN6iHA"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google-research/google-research/blob/master/gfsa/notebooks/guide_for_new_tasks.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hmLxdk2D5YQG"
      },
      "source": [
        "### Setup and imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UxOrT9g05XWV",
        "outputId": "370cbc50-ac38-475b-90f9-613e4056f1a6"
      },
      "source": [
        "!git clone https://github.com/google-research/google-research.git --depth=1"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'google-research'...\n",
            "remote: Enumerating objects: 9929, done.\u001b[K\n",
            "remote: Counting objects: 100% (9929/9929), done.\u001b[K\n",
            "remote: Compressing objects: 100% (7957/7957), done.\u001b[K\n",
            "remote: Total 9929 (delta 1409), reused 7837 (delta 1180), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (9929/9929), 109.27 MiB | 20.40 MiB/s, done.\n",
            "Resolving deltas: 100% (1409/1409), done.\n",
            "Checking out files: 100% (10428/10428), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d450jDR_5Z-q"
      },
      "source": [
        "import os\n",
        "os.chdir(\"google-research\")"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gfCDCBvN59gv",
        "outputId": "518d907f-c4e0-4c3a-96aa-b20331ee975f"
      },
      "source": [
        "!pip install flax"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting flax\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/2e/ef/1a0c6a869396e4b20a83375fe06d725b5c1711746578ee8dd6969472af41/flax-0.2.2-py3-none-any.whl (148kB)\n",
            "\r\u001b[K     |██▏                             | 10kB 16.7MB/s eta 0:00:01\r\u001b[K     |████▍                           | 20kB 19.6MB/s eta 0:00:01\r\u001b[K     |██████▋                         | 30kB 10.7MB/s eta 0:00:01\r\u001b[K     |████████▉                       | 40kB 8.7MB/s eta 0:00:01\r\u001b[K     |███████████                     | 51kB 4.3MB/s eta 0:00:01\r\u001b[K     |█████████████▎                  | 61kB 4.7MB/s eta 0:00:01\r\u001b[K     |███████████████▍                | 71kB 4.9MB/s eta 0:00:01\r\u001b[K     |█████████████████▋              | 81kB 5.3MB/s eta 0:00:01\r\u001b[K     |███████████████████▉            | 92kB 5.6MB/s eta 0:00:01\r\u001b[K     |██████████████████████          | 102kB 4.2MB/s eta 0:00:01\r\u001b[K     |████████████████████████▎       | 112kB 4.2MB/s eta 0:00:01\r\u001b[K     |██████████████████████████▌     | 122kB 4.2MB/s eta 0:00:01\r\u001b[K     |████████████████████████████▊   | 133kB 4.2MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▉ | 143kB 4.2MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 153kB 4.2MB/s \n",
            "\u001b[?25hRequirement already satisfied: msgpack in /usr/local/lib/python3.6/dist-packages (from flax) (1.0.0)\n",
            "Requirement already satisfied: matplotlib in /usr/local/lib/python3.6/dist-packages (from flax) (3.2.2)\n",
            "Requirement already satisfied: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from flax) (0.8)\n",
            "Requirement already satisfied: jax>=0.1.59 in /usr/local/lib/python3.6/dist-packages (from flax) (0.2.6)\n",
            "Requirement already satisfied: numpy>=1.12 in /usr/local/lib/python3.6/dist-packages (from flax) (1.18.5)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->flax) (1.3.1)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.6/dist-packages (from matplotlib->flax) (0.10.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->flax) (2.8.1)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.6/dist-packages (from matplotlib->flax) (2.4.7)\n",
            "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.6/dist-packages (from jax>=0.1.59->flax) (3.3.0)\n",
            "Requirement already satisfied: absl-py in /usr/local/lib/python3.6/dist-packages (from jax>=0.1.59->flax) (0.10.0)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from cycler>=0.10->matplotlib->flax) (1.15.0)\n",
            "Installing collected packages: flax\n",
            "Successfully installed flax-0.2.2\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iPo-cZD45hb_"
      },
      "source": [
        "import gast\n",
        "import numpy as np\n",
        "\n",
        "from gfsa import automaton_builder\n",
        "from gfsa import generic_ast_graphs\n",
        "from gfsa import graph_types\n",
        "from gfsa import py_ast_graphs\n",
        "from gfsa import schema_util\n",
        "from gfsa.datasets.mazes import maze_schema\n",
        "from gfsa.datasets.mazes import maze_task\n",
        "from gfsa.visualization.pprint import pprint\n",
        "from gfsa.visualization.pytrees import summarize_tree"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6GQcltBWibY"
      },
      "source": [
        "## Defining your graph domain"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kCfgkdNbalY6"
      },
      "source": [
        "The first step in using the GFSA layer for a new task is to specify how to interpret the graphs in your domain as MDPs. Specifically, you must define a set of node types, and then for each node type, define the set of possible actions (\"out edges\") the agent can take at that node, and the set of observations (\"in edges\") the agent can receive when it arrives at the node. In the codebase, this is referred to as a \"graph schema\".\n",
        "\n",
        "For the dataset of simple Python functions, the graph schema is derived from a simpler \"AST specification\":"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YkQGlO4nbepU",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "dc1e9a78-6e14-46ed-ad34-d7c70f3e1b3b"
      },
      "source": [
        "# AST specification:\n",
        "py_ast_graphs.PY_AST_SPECS"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'Add': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'And': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Assign': ASTNodeSpec(fields={'targets': FieldType.ONE_CHILD, 'value': FieldType.ONE_CHILD}, sequence_item_types={}, has_parent=True),\n",
              " 'BinOp': ASTNodeSpec(fields={'left': FieldType.ONE_CHILD, 'op': FieldType.ONE_CHILD, 'right': FieldType.ONE_CHILD}, sequence_item_types={}, has_parent=True),\n",
              " 'BoolOp': ASTNodeSpec(fields={'op': FieldType.ONE_CHILD, 'values': FieldType.NONEMPTY_SEQUENCE}, sequence_item_types={'values': 'BoolOp_values'}, has_parent=True),\n",
              " 'Break': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Call': ASTNodeSpec(fields={'func': FieldType.ONE_CHILD, 'args': FieldType.SEQUENCE, 'keywords': FieldType.NO_CHILDREN}, sequence_item_types={'args': 'Call_args'}, has_parent=True),\n",
              " 'Compare': ASTNodeSpec(fields={'left': FieldType.ONE_CHILD, 'ops': FieldType.ONE_CHILD, 'comparators': FieldType.ONE_CHILD}, sequence_item_types={}, has_parent=True),\n",
              " 'Constant': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Continue': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Div': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Eq': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Expr': ASTNodeSpec(fields={'value': FieldType.ONE_CHILD}, sequence_item_types={}, has_parent=True),\n",
              " 'For': ASTNodeSpec(fields={'target': FieldType.ONE_CHILD, 'iter': FieldType.ONE_CHILD, 'body': FieldType.NONEMPTY_SEQUENCE, 'orelse': FieldType.NO_CHILDREN}, sequence_item_types={'body': 'For_body'}, has_parent=True),\n",
              " 'FunctionDef': ASTNodeSpec(fields={'args': FieldType.ONE_CHILD, 'returns': FieldType.OPTIONAL_CHILD, 'body': FieldType.NONEMPTY_SEQUENCE, 'decorator_list': FieldType.NO_CHILDREN}, sequence_item_types={'body': 'FunctionDef_body'}, has_parent=True),\n",
              " 'Gt': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'GtE': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'If': ASTNodeSpec(fields={'test': FieldType.ONE_CHILD, 'body': FieldType.NONEMPTY_SEQUENCE, 'orelse': FieldType.SEQUENCE}, sequence_item_types={'body': 'If_body', 'orelse': 'If_orelse'}, has_parent=True),\n",
              " 'Lt': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'LtE': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Module': ASTNodeSpec(fields={'body': FieldType.NONEMPTY_SEQUENCE}, sequence_item_types={'body': 'Module_body'}, has_parent=False),\n",
              " 'Mult': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Name': ASTNodeSpec(fields={'ctx': FieldType.IGNORE}, sequence_item_types={}, has_parent=True),\n",
              " 'NotEq': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Or': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Pass': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'Return': ASTNodeSpec(fields={'value': FieldType.OPTIONAL_CHILD}, sequence_item_types={}, has_parent=True),\n",
              " 'Sub': ASTNodeSpec(fields={}, sequence_item_types={}, has_parent=True),\n",
              " 'While': ASTNodeSpec(fields={'test': FieldType.ONE_CHILD, 'body': FieldType.NONEMPTY_SEQUENCE, 'orelse': FieldType.NO_CHILDREN}, sequence_item_types={'body': 'While_body'}, has_parent=True),\n",
              " 'arguments': ASTNodeSpec(fields={'args': FieldType.SEQUENCE, 'posonlyargs': FieldType.NO_CHILDREN, 'vararg': FieldType.NO_CHILDREN, 'kwonlyargs': FieldType.NO_CHILDREN, 'kw_defaults': FieldType.NO_CHILDREN, 'kwarg': FieldType.NO_CHILDREN, 'defaults': FieldType.NO_CHILDREN}, sequence_item_types={'args': 'arguments_args'}, has_parent=True)}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 6
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2WmX-csofj3J",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "61051549-1cd3-4342-9c95-25a051770264"
      },
      "source": [
        "# Derived schema (using `generic_ast_graphs.build_ast_graph_schema`)\n",
        "py_ast_graphs.SCHEMA"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'Add': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'And': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Assign': NodeSchema(in_edges=['parent_in', 'targets_in', 'value_in'], out_edges=['parent_out', 'targets_out', 'value_out']),\n",
              " 'BinOp': NodeSchema(in_edges=['parent_in', 'left_in', 'op_in', 'right_in'], out_edges=['parent_out', 'left_out', 'op_out', 'right_out']),\n",
              " 'BoolOp': NodeSchema(in_edges=['parent_in', 'op_in', 'values_in'], out_edges=['parent_out', 'op_out', 'values_out_all', 'values_out_first', 'values_out_last']),\n",
              " 'BoolOp_values-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'Break': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Call': NodeSchema(in_edges=['parent_in', 'func_in', 'args_in', 'args_missing'], out_edges=['parent_out', 'func_out', 'args_out_all', 'args_out_first', 'args_out_last']),\n",
              " 'Call_args-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'Compare': NodeSchema(in_edges=['parent_in', 'left_in', 'ops_in', 'comparators_in'], out_edges=['parent_out', 'left_out', 'ops_out', 'comparators_out']),\n",
              " 'Constant': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Continue': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Div': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Eq': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Expr': NodeSchema(in_edges=['parent_in', 'value_in'], out_edges=['parent_out', 'value_out']),\n",
              " 'For': NodeSchema(in_edges=['parent_in', 'target_in', 'iter_in', 'body_in'], out_edges=['parent_out', 'target_out', 'iter_out', 'body_out_all', 'body_out_first', 'body_out_last']),\n",
              " 'For_body-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'FunctionDef': NodeSchema(in_edges=['parent_in', 'args_in', 'returns_in', 'returns_missing', 'body_in'], out_edges=['parent_out', 'args_out', 'returns_out', 'body_out_all', 'body_out_first', 'body_out_last']),\n",
              " 'FunctionDef_body-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'Gt': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'GtE': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'If': NodeSchema(in_edges=['parent_in', 'test_in', 'body_in', 'orelse_in', 'orelse_missing'], out_edges=['parent_out', 'test_out', 'body_out_all', 'body_out_first', 'body_out_last', 'orelse_out_all', 'orelse_out_first', 'orelse_out_last']),\n",
              " 'If_body-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'If_orelse-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'Lt': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'LtE': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Module': NodeSchema(in_edges=['body_in'], out_edges=['body_out_all', 'body_out_first', 'body_out_last']),\n",
              " 'Module_body-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'Mult': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Name': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'NotEq': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Or': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Pass': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'Return': NodeSchema(in_edges=['parent_in', 'value_in', 'value_missing'], out_edges=['parent_out', 'value_out']),\n",
              " 'Sub': NodeSchema(in_edges=['parent_in'], out_edges=['parent_out']),\n",
              " 'While': NodeSchema(in_edges=['parent_in', 'test_in', 'body_in'], out_edges=['parent_out', 'test_out', 'body_out_all', 'body_out_first', 'body_out_last']),\n",
              " 'While_body-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out']),\n",
              " 'arguments': NodeSchema(in_edges=['parent_in', 'args_in', 'args_missing'], out_edges=['parent_out', 'args_out_all', 'args_out_first', 'args_out_last']),\n",
              " 'arguments_args-seq-helper': NodeSchema(in_edges=['parent_in', 'item_in', 'next_in', 'next_missing', 'prev_in', 'prev_missing'], out_edges=['parent_out', 'item_out', 'next_out', 'prev_out'])}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bNnAv9UrbfIl"
      },
      "source": [
        "It's possible to infer the AST specification from a dataset of ASTs using `ast_spec_inference.py`. In the paper, we use two different AST specifications, one for the synthetic Python examples (shown above), and one for the Python examples written by humans (since these use many additional types of AST nodes)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFSJYg7VbzER"
      },
      "source": [
        "For the maze dataset, the node types are determined by the shape of the grid cell, and the graph schema determines which actions are valid:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sY7P6YZmcT4X",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1e847b83-2811-4163-8bc6-7eedceaecb80"
      },
      "source": [
        "maze_task.SCHEMA"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'cell_LRUD': NodeSchema(in_edges=['L_in', 'R_in', 'U_in', 'D_in'], out_edges=['L_out', 'R_out', 'U_out', 'D_out']),\n",
              " 'cell_LRUx': NodeSchema(in_edges=['L_in', 'R_in', 'U_in'], out_edges=['L_out', 'R_out', 'U_out']),\n",
              " 'cell_LRxD': NodeSchema(in_edges=['L_in', 'R_in', 'D_in'], out_edges=['L_out', 'R_out', 'D_out']),\n",
              " 'cell_LRxx': NodeSchema(in_edges=['L_in', 'R_in'], out_edges=['L_out', 'R_out']),\n",
              " 'cell_LxUD': NodeSchema(in_edges=['L_in', 'U_in', 'D_in'], out_edges=['L_out', 'U_out', 'D_out']),\n",
              " 'cell_LxUx': NodeSchema(in_edges=['L_in', 'U_in'], out_edges=['L_out', 'U_out']),\n",
              " 'cell_LxxD': NodeSchema(in_edges=['L_in', 'D_in'], out_edges=['L_out', 'D_out']),\n",
              " 'cell_xRUD': NodeSchema(in_edges=['R_in', 'U_in', 'D_in'], out_edges=['R_out', 'U_out', 'D_out']),\n",
              " 'cell_xRUx': NodeSchema(in_edges=['R_in', 'U_in'], out_edges=['R_out', 'U_out']),\n",
              " 'cell_xRxD': NodeSchema(in_edges=['R_in', 'D_in'], out_edges=['R_out', 'D_out']),\n",
              " 'cell_xxUD': NodeSchema(in_edges=['U_in', 'D_in'], out_edges=['U_out', 'D_out'])}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tC9DHd1LcUIx"
      },
      "source": [
        "As a toy example of how you might encode a new graph domain, suppose we have a network of houses connected by directed roads. Each house is adjacent to exactly one road, and each road has at least one entry and exit point but may have more.\n",
        "\n",
        "![image.png]()\n",
        "\n",
        "We can encode this structure using the following schema:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IVpdKvAscoX0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "5f2773ba-7315-4e98-e41e-ba02f57708cc"
      },
      "source": [
        "road_network_schema = {\n",
        "    # Houses are only connected to roads, so the only movement action\n",
        "    # available is to go to the road; likewise the only observation we receive\n",
        "    # after moving is the observation that we arrived from a road.\n",
        "    \"house\": graph_types.NodeSchema(\n",
        "        in_edges=[\"from_road\"],\n",
        "        out_edges=[\"to_road\"]),\n",
        "    # Roads are more complex. We can always move to a random previous or next\n",
        "    # road in the road network. We can also try to move to a house, but if there\n",
        "    # is no house, we will have to stay on the road. We denote this with a\n",
        "    # special observation (in `in_edges`).\n",
        "    \"road\": graph_types.NodeSchema(\n",
        "        in_edges=[\"from_next\", \"from_prev\", \"from_house\", \"no_house_here\"],\n",
        "        out_edges=[\"to_next\", \"to_prev\", \"to_house\"]),\n",
        "}\n",
        "road_network_schema"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'house': NodeSchema(in_edges=['from_road'], out_edges=['to_road']),\n",
              " 'road': NodeSchema(in_edges=['from_next', 'from_prev', 'from_house', 'no_house_here'], out_edges=['to_next', 'to_prev', 'to_house'])}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zwy8A7TkcxaN"
      },
      "source": [
        "## Building MDP graphs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fac7M2ZSc2Fs"
      },
      "source": [
        "Before running the GFSA layer on a specific input graph, you need to specify the result of taking each of the actions defined in the schema.\n",
        "\n",
        "For AST graphs, these transitions can be automatically computed based on the AST and its specification:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2xCp_l1GdWnh",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a3d2c72b-2855-40a5-b0cb-50e95c4808bb"
      },
      "source": [
        "the_ast = gast.parse(\"\"\"\n",
        "def test_function(foo):\n",
        "  if foo:\n",
        "    return\n",
        "  pass\n",
        "\"\"\")\n",
        "generic_ast = py_ast_graphs.py_ast_to_generic(the_ast)\n",
        "mdp_graph, id_conversion_map = generic_ast_graphs.ast_to_graph(generic_ast, ast_spec=py_ast_graphs.PY_AST_SPECS)\n",
        "\n",
        "schema_util.assert_conforms_to_schema(mdp_graph, py_ast_graphs.SCHEMA)\n",
        "\n",
        "mdp_graph"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'root__Module': GraphNode(node_type='Module', out_edges={'body_out_all': [InputTaggedNode(node_id='root_body_0__Module_body-seq-helper', in_edge='parent_in')], 'body_out_first': [InputTaggedNode(node_id='root_body_0__Module_body-seq-helper', in_edge='parent_in')], 'body_out_last': [InputTaggedNode(node_id='root_body_0__Module_body-seq-helper', in_edge='parent_in')]}),\n",
              " 'root_body_0__Module_body-seq-helper': GraphNode(node_type='Module_body-seq-helper', out_edges={'parent_out': [InputTaggedNode(node_id='root__Module', in_edge='body_in')], 'item_out': [InputTaggedNode(node_id='root_body_0_item__FunctionDef', in_edge='parent_in')], 'prev_out': [InputTaggedNode(node_id='root_body_0__Module_body-seq-helper', in_edge='prev_missing')], 'next_out': [InputTaggedNode(node_id='root_body_0__Module_body-seq-helper', in_edge='next_missing')]}),\n",
              " 'root_body_0_item__FunctionDef': GraphNode(node_type='FunctionDef', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0__Module_body-seq-helper', in_edge='item_in')], 'args_out': [InputTaggedNode(node_id='root_body_0_item_args__arguments', in_edge='parent_in')], 'returns_out': [InputTaggedNode(node_id='root_body_0_item__FunctionDef', in_edge='returns_missing')], 'body_out_all': [InputTaggedNode(node_id='root_body_0_item_body_0__FunctionDef_body-seq-helper', in_edge='parent_in'), InputTaggedNode(node_id='root_body_0_item_body_1__FunctionDef_body-seq-helper', in_edge='parent_in')], 'body_out_first': [InputTaggedNode(node_id='root_body_0_item_body_0__FunctionDef_body-seq-helper', in_edge='parent_in')], 'body_out_last': [InputTaggedNode(node_id='root_body_0_item_body_1__FunctionDef_body-seq-helper', in_edge='parent_in')]}),\n",
              " 'root_body_0_item_args__arguments': GraphNode(node_type='arguments', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item__FunctionDef', in_edge='args_in')], 'args_out_all': [InputTaggedNode(node_id='root_body_0_item_args_args_0__arguments_args-seq-helper', in_edge='parent_in')], 'args_out_first': [InputTaggedNode(node_id='root_body_0_item_args_args_0__arguments_args-seq-helper', in_edge='parent_in')], 'args_out_last': [InputTaggedNode(node_id='root_body_0_item_args_args_0__arguments_args-seq-helper', in_edge='parent_in')]}),\n",
              " 'root_body_0_item_args_args_0__arguments_args-seq-helper': GraphNode(node_type='arguments_args-seq-helper', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_args__arguments', in_edge='args_in')], 'item_out': [InputTaggedNode(node_id='root_body_0_item_args_args_0_item__Name', in_edge='parent_in')], 'prev_out': [InputTaggedNode(node_id='root_body_0_item_args_args_0__arguments_args-seq-helper', in_edge='prev_missing')], 'next_out': [InputTaggedNode(node_id='root_body_0_item_args_args_0__arguments_args-seq-helper', in_edge='next_missing')]}),\n",
              " 'root_body_0_item_args_args_0_item__Name': GraphNode(node_type='Name', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_args_args_0__arguments_args-seq-helper', in_edge='item_in')]}),\n",
              " 'root_body_0_item_body_0__FunctionDef_body-seq-helper': GraphNode(node_type='FunctionDef_body-seq-helper', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item__FunctionDef', in_edge='body_in')], 'item_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item__If', in_edge='parent_in')], 'next_out': [InputTaggedNode(node_id='root_body_0_item_body_1__FunctionDef_body-seq-helper', in_edge='prev_in')], 'prev_out': [InputTaggedNode(node_id='root_body_0_item_body_0__FunctionDef_body-seq-helper', in_edge='prev_missing')]}),\n",
              " 'root_body_0_item_body_0_item__If': GraphNode(node_type='If', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_body_0__FunctionDef_body-seq-helper', in_edge='item_in')], 'test_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item_test__Name', in_edge='parent_in')], 'body_out_all': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0__If_body-seq-helper', in_edge='parent_in')], 'body_out_first': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0__If_body-seq-helper', in_edge='parent_in')], 'body_out_last': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0__If_body-seq-helper', in_edge='parent_in')], 'orelse_out_all': [InputTaggedNode(node_id='root_body_0_item_body_0_item__If', in_edge='orelse_missing')], 'orelse_out_first': [InputTaggedNode(node_id='root_body_0_item_body_0_item__If', in_edge='orelse_missing')], 'orelse_out_last': [InputTaggedNode(node_id='root_body_0_item_body_0_item__If', in_edge='orelse_missing')]}),\n",
              " 'root_body_0_item_body_0_item_body_0__If_body-seq-helper': GraphNode(node_type='If_body-seq-helper', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item__If', in_edge='body_in')], 'item_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0_item__Return', in_edge='parent_in')], 'prev_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0__If_body-seq-helper', in_edge='prev_missing')], 'next_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0__If_body-seq-helper', in_edge='next_missing')]}),\n",
              " 'root_body_0_item_body_0_item_body_0_item__Return': GraphNode(node_type='Return', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0__If_body-seq-helper', in_edge='item_in')], 'value_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item_body_0_item__Return', in_edge='value_missing')]}),\n",
              " 'root_body_0_item_body_0_item_test__Name': GraphNode(node_type='Name', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_body_0_item__If', in_edge='test_in')]}),\n",
              " 'root_body_0_item_body_1__FunctionDef_body-seq-helper': GraphNode(node_type='FunctionDef_body-seq-helper', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item__FunctionDef', in_edge='body_in')], 'item_out': [InputTaggedNode(node_id='root_body_0_item_body_1_item__Pass', in_edge='parent_in')], 'prev_out': [InputTaggedNode(node_id='root_body_0_item_body_0__FunctionDef_body-seq-helper', in_edge='next_in')], 'next_out': [InputTaggedNode(node_id='root_body_0_item_body_1__FunctionDef_body-seq-helper', in_edge='next_missing')]}),\n",
              " 'root_body_0_item_body_1_item__Pass': GraphNode(node_type='Pass', out_edges={'parent_out': [InputTaggedNode(node_id='root_body_0_item_body_1__FunctionDef_body-seq-helper', in_edge='item_in')]})}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y7CzopTKdW4d"
      },
      "source": [
        "In the maze dataset, we precompute the destination of taking each possible action at each possible node:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t6k0w0OSdkqg",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "88ce4952-296b-4b91-be7a-983a39735263"
      },
      "source": [
        "the_maze_raw = [\n",
        "  \"███████  ████\",\n",
        "  \"████  █  █  █\",\n",
        "  \"████  ███████\",\n",
        "]\n",
        "the_maze = np.array([[c != \" \" for c in r] for r in the_maze_raw])\n",
        "mdp_graph, coordinates = maze_schema.encode_maze(the_maze)\n",
        "\n",
        "schema_util.assert_conforms_to_schema(mdp_graph, maze_task.SCHEMA)\n",
        "\n",
        "mdp_graph"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'cell_0_0': GraphNode(node_type='cell_xRxD', out_edges={'R_out': [InputTaggedNode(node_id='cell_0_1', in_edge='L_in')], 'D_out': [InputTaggedNode(node_id='cell_1_0', in_edge='U_in')]}),\n",
              " 'cell_0_1': GraphNode(node_type='cell_LRxD', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_0', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_2', in_edge='L_in')], 'D_out': [InputTaggedNode(node_id='cell_1_1', in_edge='U_in')]}),\n",
              " 'cell_0_10': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_9', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_11', in_edge='L_in')]}),\n",
              " 'cell_0_11': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_10', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_12', in_edge='L_in')]}),\n",
              " 'cell_0_12': GraphNode(node_type='cell_LxxD', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_11', in_edge='R_in')], 'D_out': [InputTaggedNode(node_id='cell_1_12', in_edge='U_in')]}),\n",
              " 'cell_0_2': GraphNode(node_type='cell_LRxD', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_1', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_3', in_edge='L_in')], 'D_out': [InputTaggedNode(node_id='cell_1_2', in_edge='U_in')]}),\n",
              " 'cell_0_3': GraphNode(node_type='cell_LRxD', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_2', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_4', in_edge='L_in')], 'D_out': [InputTaggedNode(node_id='cell_1_3', in_edge='U_in')]}),\n",
              " 'cell_0_4': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_3', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_5', in_edge='L_in')]}),\n",
              " 'cell_0_5': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_4', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_0_6', in_edge='L_in')]}),\n",
              " 'cell_0_6': GraphNode(node_type='cell_LxxD', out_edges={'L_out': [InputTaggedNode(node_id='cell_0_5', in_edge='R_in')], 'D_out': [InputTaggedNode(node_id='cell_1_6', in_edge='U_in')]}),\n",
              " 'cell_0_9': GraphNode(node_type='cell_xRxD', out_edges={'R_out': [InputTaggedNode(node_id='cell_0_10', in_edge='L_in')], 'D_out': [InputTaggedNode(node_id='cell_1_9', in_edge='U_in')]}),\n",
              " 'cell_1_0': GraphNode(node_type='cell_xRUD', out_edges={'R_out': [InputTaggedNode(node_id='cell_1_1', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_0_0', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_0', in_edge='U_in')]}),\n",
              " 'cell_1_1': GraphNode(node_type='cell_LRUD', out_edges={'L_out': [InputTaggedNode(node_id='cell_1_0', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_1_2', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_0_1', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_1', in_edge='U_in')]}),\n",
              " 'cell_1_12': GraphNode(node_type='cell_xxUD', out_edges={'U_out': [InputTaggedNode(node_id='cell_0_12', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_12', in_edge='U_in')]}),\n",
              " 'cell_1_2': GraphNode(node_type='cell_LRUD', out_edges={'L_out': [InputTaggedNode(node_id='cell_1_1', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_1_3', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_0_2', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_2', in_edge='U_in')]}),\n",
              " 'cell_1_3': GraphNode(node_type='cell_LxUD', out_edges={'L_out': [InputTaggedNode(node_id='cell_1_2', in_edge='R_in')], 'U_out': [InputTaggedNode(node_id='cell_0_3', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_3', in_edge='U_in')]}),\n",
              " 'cell_1_6': GraphNode(node_type='cell_xxUD', out_edges={'U_out': [InputTaggedNode(node_id='cell_0_6', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_6', in_edge='U_in')]}),\n",
              " 'cell_1_9': GraphNode(node_type='cell_xxUD', out_edges={'U_out': [InputTaggedNode(node_id='cell_0_9', in_edge='D_in')], 'D_out': [InputTaggedNode(node_id='cell_2_9', in_edge='U_in')]}),\n",
              " 'cell_2_0': GraphNode(node_type='cell_xRUx', out_edges={'R_out': [InputTaggedNode(node_id='cell_2_1', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_1_0', in_edge='D_in')]}),\n",
              " 'cell_2_1': GraphNode(node_type='cell_LRUx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_0', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_2', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_1_1', in_edge='D_in')]}),\n",
              " 'cell_2_10': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_9', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_11', in_edge='L_in')]}),\n",
              " 'cell_2_11': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_10', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_12', in_edge='L_in')]}),\n",
              " 'cell_2_12': GraphNode(node_type='cell_LxUx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_11', in_edge='R_in')], 'U_out': [InputTaggedNode(node_id='cell_1_12', in_edge='D_in')]}),\n",
              " 'cell_2_2': GraphNode(node_type='cell_LRUx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_1', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_3', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_1_2', in_edge='D_in')]}),\n",
              " 'cell_2_3': GraphNode(node_type='cell_LxUx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_2', in_edge='R_in')], 'U_out': [InputTaggedNode(node_id='cell_1_3', in_edge='D_in')]}),\n",
              " 'cell_2_6': GraphNode(node_type='cell_xRUx', out_edges={'R_out': [InputTaggedNode(node_id='cell_2_7', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_1_6', in_edge='D_in')]}),\n",
              " 'cell_2_7': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_6', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_8', in_edge='L_in')]}),\n",
              " 'cell_2_8': GraphNode(node_type='cell_LRxx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_7', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_9', in_edge='L_in')]}),\n",
              " 'cell_2_9': GraphNode(node_type='cell_LRUx', out_edges={'L_out': [InputTaggedNode(node_id='cell_2_8', in_edge='R_in')], 'R_out': [InputTaggedNode(node_id='cell_2_10', in_edge='L_in')], 'U_out': [InputTaggedNode(node_id='cell_1_9', in_edge='D_in')]})}"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fPlXc3f5dlJf"
      },
      "source": [
        "For the toy houses-and-roads example above, we might have a graph that looks something like this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hjgHOilYdqDa"
      },
      "source": [
        "GraphNode = graph_types.GraphNode\n",
        "InputTaggedNode = graph_types.InputTaggedNode\n",
        "mdp_graph = {\n",
        "    'R0': GraphNode(node_type='road',\n",
        "                    out_edges={\n",
        "                        'to_next': [InputTaggedNode(node_id='R1', in_edge='from_prev')],\n",
        "                        'to_prev': [InputTaggedNode(node_id='R1', in_edge='from_next')],\n",
        "                        'to_house': [InputTaggedNode(node_id='H0', in_edge='from_road')]\n",
        "                    }),\n",
        "    'R1': GraphNode(node_type='road',\n",
        "                    out_edges={\n",
        "                        'to_next': [InputTaggedNode(node_id='R0', in_edge='from_prev')],\n",
        "                        'to_prev': [InputTaggedNode(node_id='R0', in_edge='from_next'),\n",
        "                                    InputTaggedNode(node_id='R3', in_edge='from_next')],\n",
        "                        'to_house': [InputTaggedNode(node_id='H1', in_edge='from_road')]\n",
        "                    }),\n",
        "    'R2': GraphNode(node_type='road',\n",
        "                    out_edges={\n",
        "                        'to_next': [InputTaggedNode(node_id='R3', in_edge='from_prev'),\n",
        "                                    InputTaggedNode(node_id='R4', in_edge='from_prev')],\n",
        "                        'to_prev': [InputTaggedNode(node_id='R4', in_edge='from_next')],\n",
        "                        'to_house': [InputTaggedNode(node_id='R2', in_edge='no_house_here')]\n",
        "                    }),\n",
        "    'R3': GraphNode(node_type='road',\n",
        "                    out_edges={\n",
        "                        'to_next': [InputTaggedNode(node_id='R1', in_edge='from_prev')],\n",
        "                        'to_prev': [InputTaggedNode(node_id='R2', in_edge='from_next')],\n",
        "                        'to_house': [InputTaggedNode(node_id='H2', in_edge='from_road')]\n",
        "                    }),\n",
        "    'R4': GraphNode(node_type='road',\n",
        "                    out_edges={\n",
        "                        'to_next': [InputTaggedNode(node_id='R2', in_edge='from_prev')],\n",
        "                        'to_prev': [InputTaggedNode(node_id='R2', in_edge='from_next')],\n",
        "                        'to_house': [InputTaggedNode(node_id='R4', in_edge='no_house_here')]\n",
        "                    }),\n",
        "    'H0': GraphNode(node_type='house',\n",
        "                    out_edges={\n",
        "                        'to_road': [InputTaggedNode(node_id='R0', in_edge='from_house')]\n",
        "                    }),\n",
        "    'H1': GraphNode(node_type='house',\n",
        "                    out_edges={\n",
        "                        'to_road': [InputTaggedNode(node_id='R1', in_edge='from_house')]\n",
        "                    }),\n",
        "    'H2': GraphNode(node_type='house',\n",
        "                    out_edges={\n",
        "                        'to_road': [InputTaggedNode(node_id='R3', in_edge='from_house')]\n",
        "                    }),\n",
        "}\n",
        "\n",
        "schema_util.assert_conforms_to_schema(mdp_graph, road_network_schema)"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oWxKGiDHinxz"
      },
      "source": [
        "Note that every action must have at least one destination and associated observation! This is the case even for roads with no house, for which the \"to_house\" action results in staying in place and getting a special sentinel observation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LLXLYzdg62zd"
      },
      "source": [
        "## Encoding MDP graphs into GFSA-compatible NDArrays"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1sirS_gVjO-d"
      },
      "source": [
        "Once you have an MDP graph that conforms to a schema, you can use an automaton builder object to encode that graph into a set of NDArrays:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W8S-peKCjUWu",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "bf3f3861-0323-45de-ab0a-5d54c254769a"
      },
      "source": [
        "road_builder = automaton_builder.AutomatonBuilder(road_network_schema)\n",
        "pprint(summarize_tree(road_builder.encode_graph(mdp_graph, as_jax=False)))"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "(EncodedGraph(\n",
            "     initial_to_in_tagged=SparseCoordOperator(\n",
            "                              input_indices='int32(20, 1) [1, 16] nz:1.0',\n",
            "                              output_indices='int32(20, 2) [0, 17] nz:0.9',\n",
            "                              values='float32(20,) 0.9 ±0.2 [0.5, 1.0]'),\n",
            "     initial_to_special='int32(8,) [1, 6] nz:1.0',\n",
            "     in_tagged_to_in_tagged=SparseCoordOperator(\n",
            "                                input_indices='int32(54, 1) [0, 13] nz:0.94',\n",
            "                                output_indices='int32(54, 2) [0, 17] nz:0.96',\n",
            "                                values='float32(54,) 0.89 ±0.21 [0.5, 1.0]'),\n",
            "     in_tagged_to_special='int32(18,) [0, 5] nz:0.83',\n",
            "     in_tagged_node_indices='int32(18,) [0, 7] nz:0.83'),\n",
            " EncodedGraphMetadata(num_nodes=8, num_input_tagged_nodes=18))\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p-bKOQVdWo2Q"
      },
      "source": [
        "## Using the high-level Flax wrapper API\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S20Jg2VCjj3L"
      },
      "source": [
        "Once you have an encoded MDP graph, you can feed it as input to the high-level GFSA layer class, `automaton_layer.FiniteStateGraphAutomaton`.\n",
        "\n",
        "The required inputs to this layer are:\n",
        "- encoded_graph: The encoded graph for this example. This may be padded to a larger size with zeros.\n",
        "- dynamic_metadata: The true size of this example, giving the number of nodes and the number of (node, observation) pairs.\n",
        "- static_metadata: The padded size of the input, which should be the same across all graphs.\n",
        "- builder: The builder object for this automaton, which also keeps track of the schema used to encode the graph.\n",
        "\n",
        "The parameters for this layer are initialized automatically when the layer is used, according to the conventions of the (pre-Linen) `flax.deprecated.nn` API.\n",
        "\n",
        "The `variant_weights` argument is optional, and allows you to specify extra pairwise dependencies between nodes in the graph. This can be used to allow the agent to take different actions when it reaches variables that are the same as the variable used at the start location, for instance. This must sum to 1 across all variants, but it is allowed to be continuous, and can be differentiated through (so it could be the softmax-normalized output of a neural network).\n",
        "\n",
        "Useful hyperparameters:\n",
        "- `num_out_edges` determines how many types of edges to generate in parallel.\n",
        "- `num_intermediate_states` determines how many non-initial automaton hidden states there are. Higher numbers give the policy more memory but make it slower to train.\n",
        "- `share_states_across_edges` determines whether the intermediate states are shared between multiple parallel edge types. Only used if `num_out_edges > 1`.\n",
        "- `logit_scaling`: Determines whether or not to apply the learned adjustment parameters, which allow one-to-many relationships.\n",
        "\n",
        "The return value of this layer is a matrix `[num_nodes, num_nodes, num_out_edges]` containing probabilities between 0 and 1 for each pair of nodes and each edge type. Since this output is differentiable with respect to the parameters, this matrix can be used in a variety of ways. For instance, `model/end_to_end_stack.py` shows how to combine the GFSA layer with graph neural network layers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KUEURzZmW0pa"
      },
      "source": [
        "## Optional: Using the low-level GFSA building blocks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aFaKY1bjl7CG"
      },
      "source": [
        "The GFSA layer's internal operations can also be executed independently of the Flax wrapper. Some of the interesting functions include:\n",
        "\n",
        "- `automaton_builder.AutomatonBuilder.initialize_routing_params` constructs a randomly-initialized policy for a given schema\n",
        "- `automaton_builder.AutomatonBuilder.routing_softmax` normalizes a policy using softmax, to allow optimization in unconstrained logit space\n",
        "- `automaton_builder.AutomatonBuilder.build_transition_matrix` combines a policy with an encoded MDP graph to obtain a transition matrix for the absorbing Markov chain\n",
        "- `automaton_builder.all_nodes_absorbing_solve` constructs and solves the linear system to determine the absorbing distribution\n",
        "- `linear_solvers.richardson_solve` solves an arbitrary diagonally-dominant linear system using fixed-point iteration, and handles the custom gradient rules for the linear solver\n",
        "- `automaton_builder.unroll_chain_steps` unrolls the transitions of the automaton over time for visualization and analysis purposes\n",
        "- `automaton_sampling.all_nodes_particle_estimate` provides a stochastic and non-differentiable approximation to `automaton_builder.all_nodes_absorbing_solve` using multiple rollouts (used for RL ablation experiments)\n",
        "- `model_util.safe_logit` converts from probability space to logit space in order to do the one-to-many adjustment.\n",
        "\n",
        "See each function for additional documentation."
      ]
    }
  ]
}
